// ref:https://github.com/modyuan/AudioRecorder
#include "audio_recorder.h"
#include <assert.h>
const AVSampleFormat requireAudioFmt = AV_SAMPLE_FMT_FLTP;

void AudioRecorder::open()
{
    // input context

    AVDictionary *options = nullptr;
    _audioInFormatCtx = nullptr;
    int ret;

    //ref: https://ffmpeg.org/ffmpeg-devices.html
    AVInputFormat *inputFormat = nullptr;
#ifdef _WINDOWS
    inputFormat = av_find_input_format("dshow");
#elif MACOS
    if(_deviceName == "") _deviceName = ":0";
    inputFormat = av_find_input_format("avfoundation");
    //"[[VIDEO]:[AUDIO]]"
#elif UNIX
    if(_deviceName == "") _deviceName = "default";
    inputFormat=av_find_input_format("pulse");
#endif

    ret = avformat_open_input(&_audioInFormatCtx, _deviceName.c_str(), inputFormat, &options);
    if (ret != 0) {
        throw std::runtime_error("Couldn't open input audio stream.");
    }

    if (avformat_find_stream_info(_audioInFormatCtx, nullptr) < 0) {
        throw std::runtime_error("Couldn't find audio stream information.");
    }

    _audioInStream = nullptr;
    for (int i = 0; i < _audioInFormatCtx->nb_streams; i++) {
        if (_audioInFormatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            _audioInStream = _audioInFormatCtx->streams[i];
            break;
        }
    }
    if (!_audioInStream) {
        throw std::runtime_error("Couldn't find a audio stream.");
    }

    AVCodec *audioInCodec = avcodec_find_decoder(_audioInStream->codecpar->codec_id);
    _audioInCodecCtx = avcodec_alloc_context3(audioInCodec);
    avcodec_parameters_to_context(_audioInCodecCtx, _audioInStream->codecpar);

    if (avcodec_open2(_audioInCodecCtx, audioInCodec, nullptr) < 0) throw std::runtime_error("Could not open video codec.");

    // audio converter, convert other fmt to requireAudioFmt
    _audioConverter = swr_alloc_set_opts(nullptr,
                                        av_get_default_channel_layout(_audioInCodecCtx->channels),
                                        requireAudioFmt,  // aac encoder only receive this format
                                        _audioInCodecCtx->sample_rate,
                                        av_get_default_channel_layout(_audioInCodecCtx->channels),
                                        (AVSampleFormat)_audioInStream->codecpar->format,
                                        _audioInStream->codecpar->sample_rate,
                                        0, nullptr);
    swr_init(_audioConverter);

    // 2 seconds FIFO buffer
    _audioFifo = av_audio_fifo_alloc(requireAudioFmt, _audioInCodecCtx->channels,
                                    _audioInCodecCtx->sample_rate * 2);


    //--------------------------------------------------------------------------
    // output context

    // ADTS(Audio Data Transport Stream)
    ret = avformat_alloc_output_context2(&_audioOutFormatCtx, NULL, "adts", NULL);
    if (ret < 0) {
        throw std::runtime_error("Failed to alloc ouput context");
    }

    ret = avio_open(&_audioOutFormatCtx->pb, _outfile.c_str(), AVIO_FLAG_WRITE);
    if (ret < 0) {
        throw std::runtime_error("Fail to open output file.");
    }

    AVCodec *audioOutCodec = avcodec_find_encoder(AV_CODEC_ID_AAC);
    if (!audioOutCodec) throw std::runtime_error("Fail to find aac encoder. Please check your DLL.");

    _audioOutCodecCtx = avcodec_alloc_context3(audioOutCodec);
    _audioOutCodecCtx->channels = _audioInStream->codecpar->channels;
    _audioOutCodecCtx->channel_layout = av_get_default_channel_layout(_audioInStream->codecpar->channels);
    _audioOutCodecCtx->sample_rate = _audioInStream->codecpar->sample_rate;
    _audioOutCodecCtx->sample_fmt = audioOutCodec->sample_fmts[0];  //for aac , there is AV_SAMPLE_FMT_FLTP =8
    _audioOutCodecCtx->bit_rate = 32000;
    _audioOutCodecCtx->time_base.num = 1;
    _audioOutCodecCtx->time_base.den = _audioOutCodecCtx->sample_rate;

    if (_audioOutFormatCtx->oformat->flags & AVFMT_GLOBALHEADER)
        _audioOutCodecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    if (avcodec_open2(_audioOutCodecCtx, audioOutCodec, NULL) < 0)
    {
        throw std::runtime_error("Fail to open ouput audio encoder.");
    }

    //Add a new stream to output,should be called by the user before avformat_write_header() for muxing
    _audioOutStream = avformat_new_stream(_audioOutFormatCtx, audioOutCodec);
    if (_audioOutStream == NULL)
    {
        throw std::runtime_error("Fail to new a audio stream.");
    }
    avcodec_parameters_from_context(_audioOutStream->codecpar, _audioOutCodecCtx);

    // write file header
    avformat_write_header(_audioOutFormatCtx, NULL);

}

void AudioRecorder::start()
{
    _audioThread = new std::thread([this]() {
        this->_isRun = true;
        puts("Start record."); fflush(stdout);
        try {
            this->startEncode();
        }
        catch (std::exception e) {
            this->_failReason = e.what();
        }
    });
}

void AudioRecorder::stop()
{
    bool r = _isRun.exchange(false);
    if (!r) return; //avoid run twice
    _audioThread->join();

    int ret = av_write_trailer(_audioOutFormatCtx);
    if (ret < 0) throw std::runtime_error("can not write file trailer.");
    avio_close(_audioOutFormatCtx->pb);

    swr_free(&_audioConverter);
    av_audio_fifo_free(_audioFifo);

    avcodec_free_context(&_audioInCodecCtx);
    avcodec_free_context(&_audioOutCodecCtx);
    
    avformat_close_input(&_audioInFormatCtx);
	avformat_free_context(_audioOutFormatCtx);
    puts("Stop record."); 
    fflush(stdout);
}

void AudioRecorder::startEncode()
{
    AVFrame *inputFrame = av_frame_alloc();
    AVPacket *inputPacket = av_packet_alloc();

    AVPacket *outputPacket = av_packet_alloc();
    uint64_t  frameCount = 0;

    int ret;

    while (_isRun) {

        //  decoding
        ret = av_read_frame(_audioInFormatCtx, inputPacket);
        if (ret < 0) {
            throw std::runtime_error("can not read frame");
        }
        ret = avcodec_send_packet(_audioInCodecCtx, inputPacket);
        if (ret < 0) {
            throw std::runtime_error("can not send pkt in decoding");
        }
        ret = avcodec_receive_frame(_audioInCodecCtx, inputFrame);
        if (ret < 0) {
            throw std::runtime_error("can not receive frame in decoding");
        }
        //--------------------------------
        // encoding

        uint8_t **cSamples = nullptr;
        ret = av_samples_alloc_array_and_samples(&cSamples, NULL, _audioOutCodecCtx->channels, inputFrame->nb_samples, requireAudioFmt, 0);
        if (ret < 0) {
            throw std::runtime_error("Fail to alloc samples by av_samples_alloc_array_and_samples.");
        }
        ret = swr_convert(_audioConverter, cSamples, inputFrame->nb_samples, (const uint8_t**)inputFrame->extended_data, inputFrame->nb_samples);
        if (ret < 0) {
            throw std::runtime_error("Fail to swr_convert.");
        }
        if (av_audio_fifo_space(_audioFifo) < inputFrame->nb_samples) throw std::runtime_error("audio buffer is too small.");

        ret = av_audio_fifo_write(_audioFifo, (void**)cSamples, inputFrame->nb_samples);
        if (ret < 0) {
            throw std::runtime_error("Fail to write fifo");
        }

        av_freep(&cSamples[0]);

        av_frame_unref(inputFrame);
        av_packet_unref(inputPacket);

        while (av_audio_fifo_size(_audioFifo) >= _audioOutCodecCtx->frame_size) {
            AVFrame* outputFrame = av_frame_alloc();
            outputFrame->nb_samples = _audioOutCodecCtx->frame_size;
            outputFrame->channels = _audioInCodecCtx->channels;
            outputFrame->channel_layout = av_get_default_channel_layout(_audioInCodecCtx->channels);
            outputFrame->format = requireAudioFmt;
            outputFrame->sample_rate = _audioOutCodecCtx->sample_rate;

            ret = av_frame_get_buffer(outputFrame, 0);
            assert(ret >= 0);
            ret = av_audio_fifo_read(_audioFifo, (void**)outputFrame->data, _audioOutCodecCtx->frame_size);
            assert(ret >= 0);

            outputFrame->pts = frameCount * _audioOutStream->time_base.den * 1024 / _audioOutCodecCtx->sample_rate;

            ret = avcodec_send_frame(_audioOutCodecCtx, outputFrame);
            if (ret < 0) {
                throw std::runtime_error("Fail to send frame in encoding");
            }
            av_frame_free(&outputFrame);
            ret = avcodec_receive_packet(_audioOutCodecCtx, outputPacket);
            if (ret == AVERROR(EAGAIN)) {
                continue;
            }
            else if (ret < 0) {
                throw std::runtime_error("Fail to receive packet in encoding");
            }


            outputPacket->stream_index = _audioOutStream->index;
            outputPacket->duration = _audioOutStream->time_base.den * 1024 / _audioOutCodecCtx->sample_rate;
            outputPacket->dts = outputPacket->pts =frameCount * _audioOutStream->time_base.den * 1024 / _audioOutCodecCtx->sample_rate;

            frameCount++;

            ret = av_write_frame(_audioOutFormatCtx, outputPacket);
            av_packet_unref(outputPacket);

        }

    }

    av_packet_free(&inputPacket);
    av_packet_free(&outputPacket);
    av_frame_free(&inputFrame);
    printf("encode %lld audio packets in total.\n", frameCount);
}
